from xml.dom.minidom import parse, Text
from matplotlib.patches import Rectangle
import torch
import torchvision
import matplotlib.pyplot as plt
import numpy as np

__all__ = ["draw_single_image"]


def draw_single_image(image, targets, mode, show=True, inches_size=(10, 10)):
    # image 处理
    if isinstance(image, np.ndarray):
        if image.dtype == np.float32:
            image = (image * 255).astype(np.uint8)
    elif isinstance(image, torch.Tensor):
        image *= 255
        image = image.permute(1, 2, 0).numpy().astype(np.uint8)
        # box处理
    if isinstance(targets, torch.Tensor):
        boxes = targets
    elif isinstance(targets, dict):
        boxes = targets["boxes"]
    else:
        raise Exception("targets error")
    H, W = image.shape[:2]
    boxes[:, 0::2] *= W
    boxes[:, 1::2] *= H
    ax: plt.Axes = plt.gca()
    fig: plt.Figure = plt.gcf()
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)
    fig.set_size_inches(inches_size)
    if mode != "xywh":
        boxes = torchvision.ops.box_convert(boxes, mode, "xywh").tolist()
    for box in boxes:
        ax.add_patch(Rectangle((box[0], box[1]), box[2], box[3], fill=None))
    plt.imshow(image)
    if show:
        plt.show()


def xml2dict(xml):
    def dfs_parse(dom, data):
        for child in dom.childNodes:
            if len(child.childNodes) == 1:  # 文本节点
                data[child.tagName] = child.firstChild.nodeValue
            elif not isinstance(child, Text):
                ans = {}
                if child.tagName not in data:
                    data[child.tagName] = ans
                elif isinstance(data[child.tagName], dict):
                    data[child.tagName] = [data[child.tagName], ans]
                else:
                    data[child.tagName].append(ans)
                dfs_parse(child, ans)
        return data

    return dfs_parse(parse(xml), {})
